# -*- coding: utf-8 -*-

import os
import datetime
import pandas as pd
import logging
import tushare as ts

from pyquant import Strategy, SimulatedExecutionHandler, BasicPortfolio, Backtest, MarketOrder, LimitOrder


class MyStrategy(Strategy):

    def init(self, context):
        """
        init 必须实现,
        必须初始化stocks数组, 初始化之后引擎会根据数组发送bar数据.
        """
        context.stocks = ['600008', '600018']
        context.long_window = 10
        context.short_window = 5
        context.bought = {}
        context.bought['600008'] = False
        context.bought['600018'] = False
        context.count = 1
        self.new_stocks = []

    def before_trading(self, context):
        # 每30天选择pe值最小的10只股票, 在handle bar中买入
        context.count = context.count + 1
        if context.count % 30 == 0:
            now = context.data_handler.now
            tdata = ts.get_stock_basics(date=now.strftime('%Y-%m-%d'))
            tdata = tdata[tdata['pe'] > 0]
            context.new_stocks = tdata.sort_values(by='pe', ascending=True)[:10].index.tolist()
            # stocks 变动需要手工更新股票池, 暂时不支持交易
            context.data_handler.update_universe(context.new_stocks)

    def handle_bar(self, context, bar_dict):
        """
        买入新的股票, 旧的股票卖出
        """
        if context.count % 30 != 0:
            return
        for s in context.new_stocks:
            if s not in context.stocks:
                context.stocks.append(s)
        for s in context.stocks:
            if s not in context.new_stocks:
                self.order_target_percent(s, 0)
            else:
                # 每只股票持有10%的资金
                self.order_target_percent(s, 0.1)



if __name__ == '__main__':
    csv_dir = os.path.join(os.getcwd(), 'testdata')
    # print(csv_dir)
    symbol_list = ['600008', '600018']
    initial_capital = 100000.0
    heartbeat = 0.0
    start_date = datetime.datetime(2016, 7, 2, 0, 0)
    end_date = datetime.datetime(2016, 12, 30, 23, 59)

    backtest = Backtest(initial_capital,
                        SimulatedExecutionHandler,
                        BasicPortfolio, MyStrategy,
                        slippage_type='fixed', commission_type='default', csv_dir=csv_dir, start_date=start_date, end_date=end_date)

    backtest.setup_logger(logging.DEBUG)
    positions, holdings = backtest.run_backtest()
    print(holdings.tail())
    # 画图
    backtest.plot()
